package com.ww.springboot.es.utils;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.WrapperQueryBuilder;
import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.aggregations.AggregatorFactories;
import org.elasticsearch.search.aggregations.BaseAggregationBuilder;
import org.elasticsearch.search.aggregations.BucketOrder;
import org.elasticsearch.search.aggregations.bucket.nested.NestedAggregationBuilder;
import org.elasticsearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.sort.SortBuilders;
import org.elasticsearch.search.sort.SortOrder;
import org.springframework.data.domain.PageRequest;
import org.springframework.data.elasticsearch.core.query.*;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import static org.elasticsearch.index.query.QueryBuilders.wrapperQuery;

@Component
public class ElasticsearchJsonUtil {

    static final Integer INDEX_MAX_RESULT_WINDOW = 10_000;

    private static final ObjectMapper objectMapper = new ObjectMapper();

    public static NativeSearchQueryBuilder fromJson(String jsonQuery) {
        try {
            JsonNode rootNode = objectMapper.readTree(jsonQuery);
            NativeSearchQueryBuilder builder = new NativeSearchQueryBuilder();

            // 处理 query 部分
            JsonNode queryNode = rootNode.get("query");
            if (queryNode != null) {
                //Wrapper 查询构造器的 toString 的内容和常规的 DSL 语言不一致
                WrapperQueryBuilder wrapperQuery = QueryBuilders.wrapperQuery(queryNode.toString());
                builder.withQuery(wrapperQuery);
            }
            // 处理聚合部分
            processAggregations(rootNode, builder);

            // 处理分页
            //processPagination(rootNode, builder);

            // 处理排序
            processSort(rootNode, builder);

            // 处理_source
            processSource(rootNode, builder);
            // 处理其他参数...

            return builder;
        } catch (Exception e) {
            throw new RuntimeException("Failed to parse JSON query", e);
        }
    }

    // 在 processSource 方法后添加以下方法
    private static void processAggregations(JsonNode rootNode, NativeSearchQueryBuilder builder) {
        // 这里需要根据聚合类型进行处理
        // 由于聚合结构复杂，可以通过 WrapperAggregationBuilder 来处理
        TermsAggregationBuilder builder1 = AggregationBuilders.terms("by_pms_sku")
                .field("packageGoods.productQuantity")
                .size(100)
                .order(BucketOrder.compound(
                        BucketOrder.aggregation("total_quantity", false), // 首先按total_quantity降序
                        BucketOrder.count(true)));

        builder1.subAggregation(AggregationBuilders
                .sum("total_quantity")
                .field("packageGoods.productQuantity"));

        NestedAggregationBuilder aggregationBuilder = AggregationBuilders
                .nested("package_goods","packageGoods")
                .subAggregation(builder1);

        builder.addAggregation(aggregationBuilder);
    }


    private static void processPagination(JsonNode rootNode, NativeSearchQueryBuilder builder) {
        int from = rootNode.path("from").asInt(0);
        int size = rootNode.path("size").asInt(10);
        if (size > 0) {
            builder.withPageable(PageRequest.of(Math.max(0, from / size), size));
        }
    }

    private static void processSort(JsonNode rootNode, NativeSearchQueryBuilder builder) {
        JsonNode sortNode = rootNode.get("sort");
        if (sortNode != null && sortNode.isArray()) {
            sortNode.forEach(sortField -> {
                if (sortField.isObject()) {
                    sortField.fields().forEachRemaining(entry -> {
                        processSortField(entry.getKey(), entry.getValue(), builder);
                    });
                }
            });
        }
    }

    private static void processSortField(String fieldName, JsonNode orderNode, NativeSearchQueryBuilder builder) {
        if (orderNode.isTextual()) {
            String order = orderNode.asText();
            SortOrder sortOrder = "desc".equalsIgnoreCase(order) ? SortOrder.DESC : SortOrder.ASC;
            builder.withSort(SortBuilders.fieldSort(fieldName).order(sortOrder));
        } else if (orderNode.isObject()) {
            String order = orderNode.path("order").asText("asc");
            SortOrder sortOrder = "desc".equalsIgnoreCase(order) ? SortOrder.DESC : SortOrder.ASC;
            builder.withSort(SortBuilders.fieldSort(fieldName).order(sortOrder));
        }
    }

    // 在你的 ElasticsearchJsonUtil 类中添加以下方法
    private static void processSource(JsonNode rootNode, NativeSearchQueryBuilder builder) {
        JsonNode sourceNode = rootNode.get("_source");
        if (sourceNode != null) {
            if (sourceNode.isBoolean()) {
                // 处理 {"_source": false} 的情况
                if (!sourceNode.asBoolean()) {
                    builder.withSourceFilter(new FetchSourceFilter(new String[]{}, null));
                }
            } else if (sourceNode.isArray()) {
                // 处理 {"_source": ["field1", "field2"]} 的情况
                List<String> includes = new ArrayList<>();
                sourceNode.forEach(node -> includes.add(node.asText()));
                builder.withSourceFilter(new FetchSourceFilter(includes.toArray(new String[0]), null));
            } else if (sourceNode.isObject()) {
                // 处理 {"_source": {"includes": [...], "excludes": [...]}} 的情况
                JsonNode includesNode = sourceNode.get("includes");
                JsonNode excludesNode = sourceNode.get("excludes");

                String[] includes = null;
                String[] excludes = null;

                if (includesNode != null && includesNode.isArray()) {
                    List<String> includesList = new ArrayList<>();
                    includesNode.forEach(node -> includesList.add(node.asText()));
                    includes = includesList.toArray(new String[0]);
                }

                if (excludesNode != null && excludesNode.isArray()) {
                    List<String> excludesList = new ArrayList<>();
                    excludesNode.forEach(node -> excludesList.add(node.asText()));
                    excludes = excludesList.toArray(new String[0]);
                }

                builder.withSourceFilter(new FetchSourceFilter(includes, excludes));
            }
        }
    }


    public static SearchRequest getSearchRequest(String queryJson, String indexName) {
        NativeSearchQueryBuilder queryBuilder = ElasticsearchJsonUtil.fromJson(queryJson);
        NativeSearchQuery query = queryBuilder.build();
        SearchRequest searchRequest = ElasticsearchJsonUtil.prepareSearchRequest(query, indexName);
        QueryBuilder elasticsearchQuery = getQuery(query);
        QueryBuilder elasticsearchFilter = getFilter(query);

        if (elasticsearchQuery != null) {
            searchRequest.source().query(elasticsearchQuery);
        } else {
            searchRequest.source().query(QueryBuilders.matchAllQuery());
        }

        if (elasticsearchFilter != null) {
            searchRequest.source().postFilter(elasticsearchFilter);
        }
        return searchRequest;
    }

    public static SearchRequest getSearchRequest(NativeSearchQuery query, String indexName) {
        SearchRequest searchRequest = ElasticsearchJsonUtil.prepareSearchRequest(query, indexName);
        QueryBuilder elasticsearchQuery = getQuery(query);
        QueryBuilder elasticsearchFilter = getFilter(query);

        if (elasticsearchQuery != null) {
            searchRequest.source().query(elasticsearchQuery);
        } else {
            searchRequest.source().query(QueryBuilders.matchAllQuery());
        }

        if (elasticsearchFilter != null) {
            searchRequest.source().postFilter(elasticsearchFilter);
        }
        return searchRequest;
    }

    public static SearchRequest prepareSearchRequest(Query query, String indexName) {
        SearchRequest request = new SearchRequest(indexName);
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
        sourceBuilder.version(true);
        sourceBuilder.trackScores(query.getTrackScores());

        if (query.getSourceFilter() != null) {
            SourceFilter sourceFilter = query.getSourceFilter();
            sourceBuilder.fetchSource(sourceFilter.getIncludes(), sourceFilter.getExcludes());
        }

        if (query.getPageable().isPaged()) {
            sourceBuilder.from((int) query.getPageable().getOffset());
            sourceBuilder.size(query.getPageable().getPageSize());
        } else {
            sourceBuilder.from(0);
            sourceBuilder.size(INDEX_MAX_RESULT_WINDOW);
        }

        if (!query.getFields().isEmpty()) {
            sourceBuilder.fetchSource(query.getFields().toArray(new String[0]), null);
        }

        if (query.getIndicesOptions() != null) {
            request.indicesOptions(query.getIndicesOptions());
        }

        if (query.isLimiting()) {
            sourceBuilder.size(query.getMaxResults());
        }

        if (query.getMinScore() > 0) {
            sourceBuilder.minScore(query.getMinScore());
        }

        if (query.getPreference() != null) {
            request.preference(query.getPreference());
        }

        if (query.getSearchType() != null) {
            request.searchType(query.getSearchType());
        }

        if (query.getTrackTotalHits()) {
            sourceBuilder.trackTotalHits(query.getTrackTotalHits());
        }

        if (StringUtils.hasLength(query.getRoute())) {
            request.routing(query.getRoute());
        }

        request.source(sourceBuilder);
        return request;
    }

    private static QueryBuilder getQuery(Query query) {
        QueryBuilder elasticsearchQuery;

        if (query instanceof NativeSearchQuery) {
            NativeSearchQuery searchQuery = (NativeSearchQuery) query;
            elasticsearchQuery = searchQuery.getQuery();
        } else if (query instanceof StringQuery) {
            StringQuery stringQuery = (StringQuery) query;
            elasticsearchQuery = wrapperQuery(stringQuery.getSource());
        } else {
            throw new IllegalArgumentException("unhandled Query implementation " + query.getClass().getName());
        }

        return elasticsearchQuery;
    }

    private static QueryBuilder getFilter(Query query) {
        QueryBuilder elasticsearchFilter;

        if (query instanceof NativeSearchQuery) {
            NativeSearchQuery searchQuery = (NativeSearchQuery) query;
            elasticsearchFilter = searchQuery.getFilter();
        } else if (query instanceof StringQuery) {
            elasticsearchFilter = null;
        } else {
            throw new IllegalArgumentException("unhandled Query implementation " + query.getClass().getName());
        }

        return elasticsearchFilter;
    }
}
